﻿using System;
using System.Buffers;
using System.Diagnostics;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
using PooledGrowableBufferHelper;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;
using SixLabors.ImageSharp.Processing;
using TiffLibrary;
using WholeSlideImageLibrary.DeepZoom;

namespace WholeSlideImageLibrary.Formats.GenericTiff
{
    internal sealed class GenericTiffDeepZoomTileSource : DeepZoomTileSource
    {
        private readonly LayerDecoder[] _layers;
        private readonly int _width;
        private readonly int _height;
        private readonly int _tileSize;


        public GenericTiffDeepZoomTileSource(TiffFileContentSource contentSource, GenericTiffLevelDecoder[] levels, int tileSize)
        {
            GenericTiffLevelDecoder level0 = levels[0];

            // 构建Layers数组
            DeepZoomLayer[]? dzLayers = DeepZoomHelper.CalculateDeepZoomLayers(level0.Width, level0.Height);
            LayerDecoder[] layers = new LayerDecoder[dzLayers.Length];

            for (int i = 0; i < layers.Length; i++)
            {
                DeepZoomLayer dzLayer = dzLayers[i];
                int layerDownsample = 1 << (dzLayers.Length - i - 1);

                // 找到合适的Level
                GenericTiffLevelDecoder? suitableLevel = null;
                if (layerDownsample <= level0.Downsample)
                {
                    suitableLevel = level0;
                }
                if (suitableLevel is null)
                {
                    for (int j = 1; j < levels.Length; j++)
                    {
                        if (layerDownsample < levels[j].Downsample)
                        {
                            suitableLevel = levels[j - 1];
                            break;
                        }
                    }
                }
                if (suitableLevel is null)
                {
                    suitableLevel = levels[levels.Length - 1];
                }

                // 检查是否可用fast path
                if (dzLayer.Width == suitableLevel.Width && dzLayer.Height == suitableLevel.Height)
                {
                    if (suitableLevel.SupportFastJpegReading)
                    {
                        layers[i] = new FastJpegLayerDecoder(suitableLevel.ImageDecoder, contentSource, suitableLevel.Offsets!, suitableLevel.ByteCounts!, suitableLevel.JpegTables!)
                        {
                            TileSize = tileSize,
                            LayerWidth = dzLayer.Width,
                            LayerHeight = dzLayer.Height
                        };
                    }
                    else
                    {
                        layers[i] = new NormalLayerDecoder(suitableLevel.ImageDecoder)
                        {
                            TileSize = tileSize,
                            LayerWidth = dzLayer.Width,
                            LayerHeight = dzLayer.Height
                        };
                    }
                }
                else
                {
                    layers[i] = new ScaledLayerDecoder(suitableLevel.ImageDecoder)
                    {
                        TileSize = tileSize,
                        LayerWidth = dzLayer.Width,
                        LayerHeight = dzLayer.Height,
                        LayerDownsample = layerDownsample,
                        Level0Downsample = suitableLevel.Downsample
                    };
                }
            }

            _layers = layers;
            _width = level0.Width;
            _height = level0.Height;
            _tileSize = tileSize;
        }

        public override EncodedImageFileFormat TileFileFormat => EncodedImageFileFormat.Jpeg;

        public override string GetDefinitionFile()
        {
            return "<?xml version=\"1.0\" encoding=\"UTF-8\"?><Image xmlns=\"http://schemas.microsoft.com/deepzoom/2008\" Format=\"jpg\" Overlap=\"0\" TileSize=\"" + _tileSize + "\"><Size Height=\"" + _height + "\" Width=\"" + _width + "\" /></Image>";
        }

        public override bool CheckTileExists(int layer, int column, int row)
        {
            LayerDecoder[]? layers = _layers;
            if (layers is null)
            {
                throw new InvalidOperationException(nameof(GenericTiffDeepZoomTileSource) + " is not properly initialized.");
            }
            if ((uint)layer >= (uint)layers.Length)
            {
                return false;
            }

            return layers[layer].CheckTileExists(row, column);
        }

        public override Task WriteToBufferAsync(int layer, int column, int row, IBufferWriter<byte> writer, CancellationToken cancellationToken)
        {
            LayerDecoder[]? layers = _layers;
            if (layers is null)
            {
                throw new InvalidOperationException(nameof(GenericTiffDeepZoomTileSource) + " is not properly initialized.");
            }
            if ((uint)layer >= (uint)layers.Length)
            {
                throw new ArgumentOutOfRangeException(nameof(layer));
            }

            return layers[layer].WriteTileAsync(row, column, writer, cancellationToken);
        }

        public override Task WriteToStreamAsync(int layer, int column, int row, Stream stream, CancellationToken cancellationToken)
        {
            LayerDecoder[]? layers = _layers;
            if (layers is null)
            {
                throw new InvalidOperationException(nameof(GenericTiffDeepZoomTileSource) + " is not properly initialized.");
            }
            if ((uint)layer >= (uint)layers.Length)
            {
                throw new ArgumentOutOfRangeException(nameof(layer));
            }

            return layers[layer].WriteTileAsync(row, column, stream, cancellationToken);
        }

        abstract class LayerDecoder
        {
            public LayerDecoder(TiffImageDecoder decoder)
            {
                LevelDecoder = decoder;
            }

            public TiffImageDecoder LevelDecoder { get; }
            public int TileSize { get; set; }
            public int LayerWidth { get; set; }
            public int LayerHeight { get; set; }

            public bool CheckTileExists(int rowIndex, int columnIndex)
            {
                int tileSize = TileSize;
                int horizontalTileCount = (LayerWidth + tileSize - 1) / tileSize;
                int verticalTileCount = (LayerHeight + tileSize - 1) / tileSize;
                if ((uint)columnIndex >= (uint)horizontalTileCount)
                {
                    return false;
                }
                if ((uint)rowIndex >= (uint)verticalTileCount)
                {
                    return false;
                }
                return true;
            }

            public abstract Task WriteTileAsync(int rowIndex, int columnIndex, Stream stream, CancellationToken cancellationToken);

            public abstract Task WriteTileAsync(int rowIndex, int columnIndex, IBufferWriter<byte> writer, CancellationToken cancellationToken);
        }

        class NormalLayerDecoder : LayerDecoder
        {
            public NormalLayerDecoder(TiffImageDecoder decoder) : base(decoder) { }

            protected async Task WriteTileNormalAsync(int rowIndex, int columnIndex, Stream stream, CancellationToken cancellationToken)
            {
                // 检查参数
                int tileSize = TileSize;
                int horizontalTileCount = (LayerWidth + tileSize - 1) / tileSize;
                int verticalTileCount = (LayerHeight + tileSize - 1) / tileSize;
                if ((uint)columnIndex >= (uint)horizontalTileCount)
                {
                    throw new ArgumentOutOfRangeException(nameof(rowIndex));
                }
                if ((uint)rowIndex >= (uint)verticalTileCount)
                {
                    throw new ArgumentOutOfRangeException(nameof(columnIndex));
                }

                int readWidth = Math.Min(tileSize, LayerWidth - columnIndex * tileSize);
                int readHeight = Math.Min(tileSize, LayerHeight - rowIndex * tileSize);

                // 解码后重新编码
                using var image = new Image<Rgba32>(readWidth, readHeight);
                await LevelDecoder.DecodeAsync(new TiffPoint(columnIndex * tileSize, rowIndex * tileSize), image, cancellationToken).ConfigureAwait(false);
                using PooledMemoryStream ms = PooledMemoryStreamManager.Shared.GetStream();
                image.SaveAsJpeg(stream);
            }

            protected async Task WriteTileNormalAsync(int rowIndex, int columnIndex, IBufferWriter<byte> writer, CancellationToken cancellationToken)
            {
                using PooledMemoryStream ms = PooledMemoryStreamManager.Shared.GetStream();
                await WriteTileNormalAsync(rowIndex, columnIndex, (Stream)ms, cancellationToken).ConfigureAwait(false);
                Copy(ms.ToReadOnlySequence(), writer);
            }

            public override Task WriteTileAsync(int rowIndex, int columnIndex, Stream stream, CancellationToken cancellationToken)
                => WriteTileNormalAsync(rowIndex, columnIndex, stream, cancellationToken);

            public override Task WriteTileAsync(int rowIndex, int columnIndex, IBufferWriter<byte> writer, CancellationToken cancellationToken)
                => WriteTileNormalAsync(rowIndex, columnIndex, writer, cancellationToken);
        }

        class ScaledLayerDecoder : LayerDecoder
        {
            public ScaledLayerDecoder(TiffImageDecoder decoder) : base(decoder) { }

            public double Level0Downsample { get; set; }
            public int LayerDownsample { get; set; }

            public override async Task WriteTileAsync(int rowIndex, int columnIndex, Stream stream, CancellationToken cancellationToken)
            {
                // 检查参数
                int tileSize = TileSize;
                int horizontalTileCount = (LayerWidth + tileSize - 1) / tileSize;
                int verticalTileCount = (LayerHeight + tileSize - 1) / tileSize;
                if ((uint)columnIndex >= (uint)horizontalTileCount)
                {
                    throw new ArgumentOutOfRangeException(nameof(rowIndex));
                }
                if ((uint)rowIndex >= (uint)verticalTileCount)
                {
                    throw new ArgumentOutOfRangeException(nameof(columnIndex));
                }

                // DeepZoom层的读取参数
                int layerXOffset = columnIndex * tileSize;
                int layerYOffset = rowIndex * tileSize;
                int layerReadWidth = Math.Min(tileSize, LayerWidth - columnIndex * tileSize);
                int layerReadHeight = Math.Min(tileSize, LayerHeight - rowIndex * tileSize);

                // 转换成TIFF中层的参数
                TiffImageDecoder level = LevelDecoder;
                double downsample = LayerDownsample / Level0Downsample;
                int levelReadWidth = (int)(layerReadWidth * downsample);
                int levelReadHeight = (int)(layerReadHeight * downsample);
                int levelXOffset = (int)(layerXOffset * downsample);
                int levelYOffset = (int)(layerYOffset * downsample);

                // 读取
                using var tile = new Image<Rgba32>(levelReadWidth, levelReadHeight);
                await level.DecodeAsync(new TiffPoint(levelXOffset, levelYOffset), tile, cancellationToken).ConfigureAwait(false);

                // 缩放回DeepZoom层
                tile.Mutate(ctx => ctx.Resize(new Size(layerReadWidth, layerReadHeight)));

                // 编码
                using PooledMemoryStream ms = PooledMemoryStreamManager.Shared.GetStream();
                tile.SaveAsJpeg(stream);
            }

            public override async Task WriteTileAsync(int rowIndex, int columnIndex, IBufferWriter<byte> writer, CancellationToken cancellationToken)
            {
                using PooledMemoryStream ms = PooledMemoryStreamManager.Shared.GetStream();
                await WriteTileAsync(rowIndex, columnIndex, (Stream)ms, cancellationToken).ConfigureAwait(false);
                Copy(ms.ToReadOnlySequence(), writer);
            }
        }

        class FastJpegLayerDecoder : NormalLayerDecoder
        {
            public FastJpegLayerDecoder(TiffImageDecoder decoder, TiffFileContentSource fileContentSource, long[] offsets, int[] byteCounts, byte[] jpegTables) : base(decoder)
            {
                FileContentSource = fileContentSource;
                Offsets = offsets;
                ByteCounts = byteCounts;
                JpegTables = jpegTables;
            }

            public TiffFileContentSource FileContentSource { get; }
            public long[] Offsets { get; }
            public int[] ByteCounts { get; }
            public byte[] JpegTables { get; }

            public override async Task WriteTileAsync(int rowIndex, int columnIndex, IBufferWriter<byte> writer, CancellationToken cancellationToken)
            {
                // 检查参数
                int horizontalTileCount = (LayerWidth + TileSize - 1) / TileSize;
                int verticalTileCount = (LayerHeight + TileSize - 1) / TileSize;
                if ((uint)columnIndex >= (uint)horizontalTileCount)
                {
                    throw new ArgumentOutOfRangeException(nameof(rowIndex));
                }
                if ((uint)rowIndex >= (uint)verticalTileCount)
                {
                    throw new ArgumentOutOfRangeException(nameof(columnIndex));
                }

                if (columnIndex == (horizontalTileCount - 1) || (rowIndex == (verticalTileCount - 1)))
                {
                    // 因为TIFF规定Tile的最小大小（存在padding），但DeepZoom格式没有规定
                    // 为了让最外层的Tile大小符合这一层的真实大小
                    // 最外层的Tile使用常规流程(NormalLayerDecoder)来读取
                    await WriteTileNormalAsync(rowIndex, columnIndex, writer, cancellationToken).ConfigureAwait(false);
                    return;
                }

                int index = rowIndex * horizontalTileCount + columnIndex;
                int tileByteCount = ByteCounts[index];
                if (tileByteCount == 0)
                {
                    // 稀疏图像
                    await WriteTileNormalAsync(rowIndex, columnIndex, writer, cancellationToken).ConfigureAwait(false);
                    return;
                }

                TiffFileContentReader reader = await FileContentSource.OpenReaderAsync(cancellationToken).ConfigureAwait(false);
                try
                {
                    byte[] jpegTables = JpegTables;
                    if (jpegTables is null || jpegTables.Length == 0)
                    {
                        // 没有JPEGTables字段，直接将TIFF文件中的JPEG流复制到输出
                        Memory<byte> buffer = writer.GetMemory(tileByteCount);
                        int bytesRead = await reader.ReadAsync(Offsets[index], buffer.Slice(0, tileByteCount), cancellationToken).ConfigureAwait(false);
                        if (bytesRead != tileByteCount)
                        {
                            throw new InvalidDataException();
                        }
                        writer.Advance(tileByteCount);
                    }
                    else
                    {
                        // 有JPEGTables字段，需要将该字段和TIFF文件中的JPEG流进行整合
                        // 我们这里手动添加SOI Marker，并将原来JPEG流中的SOI Marker替换成0xFF(padding)
                        Memory<byte> buffer = writer.GetMemory(tileByteCount + jpegTables.Length + 2);

                        // SOI
                        WriteStartOfFrameMarker(buffer.Span.Slice(0, 2));
                        // JPEGTables
                        CopyJpegTables(jpegTables, buffer.Span.Slice(2));
                        if (jpegTables.Length >= 2)
                        {
                            ReplaceStartOfFrameMarker(buffer.Span.Slice(2, 2));
                            ReplaceEndOfFrameMarker(buffer.Span.Slice(jpegTables.Length, 2));
                        }
                        // JPEG stream
                        int bytesRead = await reader.ReadAsync(Offsets[index], buffer.Slice(2 + jpegTables.Length, tileByteCount), cancellationToken).ConfigureAwait(false);
                        if (bytesRead != tileByteCount)
                        {
                            throw new InvalidDataException();
                        }
                        if (bytesRead > 2)
                        {
                            ReplaceStartOfFrameMarker(buffer.Span.Slice(2 + jpegTables.Length, 2));
                        }
                        writer.Advance(tileByteCount + jpegTables.Length + 2);
                    }
                }
                finally
                {
                    await reader.DisposeAsync().ConfigureAwait(false);
                }
            }

            public override Task WriteTileAsync(int rowIndex, int columnIndex, Stream stream, CancellationToken cancellationToken)
            {
                return WriteTileAsync(rowIndex, columnIndex, new StreamWrapper(stream), cancellationToken);
            }

            private static void WriteStartOfFrameMarker(Span<byte> buffer)
            {
                Debug.Assert(buffer.Length == 2);
                buffer[1] = 0xD8;
                buffer[0] = 0xFF;
            }

            private static void CopyJpegTables(byte[] jpegTables, Span<byte> buffer)
            {
                Debug.Assert(buffer.Length >= jpegTables.Length);
                jpegTables.AsSpan().CopyTo(buffer);
            }

            private static void ReplaceStartOfFrameMarker(Span<byte> buffer)
            {
                Debug.Assert(buffer.Length == 2);
                if (buffer[1] == 0xD8 && buffer[0] == 0xFF)
                {
                    buffer[1] = 0xFF;
                }
            }

            private static void ReplaceEndOfFrameMarker(Span<byte> buffer)
            {
                Debug.Assert(buffer.Length == 2);
                if (buffer[1] == 0xD9 && buffer[0] == 0xFF)
                {
                    buffer[1] = 0xFF;
                }
            }
        }

        private static void Copy(ReadOnlySequence<byte> source, IBufferWriter<byte> destination)
        {
            while (!source.IsEmpty)
            {
                Span<byte> buffer = destination.GetSpan();
                int copySize = (int)Math.Min(source.Length, buffer.Length);
                source.Slice(0, copySize).CopyTo(buffer);
                source = source.Slice(copySize);
                destination.Advance(copySize);
            }
        }
    }
}
